//
// Created by Cybaster on 2024/3/30.
//

#include "helper.h"

#include <vector>

#include <fstream>

namespace stitch {


    cv::Mat stitch(const cv::Mat &base, const cv::Mat &attach, const cv::Matx33f &H) {
        cv::Mat corners(1, 4, CV_32FC2);
        corners.at<cv::Vec2f>(0, 0) = cv::Vec2f(0, 0);
        corners.at<cv::Vec2f>(0, 1) = cv::Vec2f(0, attach.rows);
        corners.at<cv::Vec2f>(0, 2) = cv::Vec2f(attach.cols, 0);
        corners.at<cv::Vec2f>(0, 3) = cv::Vec2f(attach.cols, attach.rows);
        perspectiveTransform(corners, corners, H);

        float x_start = std::min(std::min(corners.at<cv::Vec2f>(0, 0)[0], corners.at<cv::Vec2f>(0, 1)[0]), (float) 0);
        float x_end = std::max(std::max(corners.at<cv::Vec2f>(0, 2)[0], corners.at<cv::Vec2f>(0, 3)[0]),
                               (float) base.cols);
        float y_start = std::min(std::min(corners.at<cv::Vec2f>(0, 0)[1], corners.at<cv::Vec2f>(0, 2)[1]), (float) 0);
        float y_end = std::max(std::max(corners.at<cv::Vec2f>(0, 1)[1], corners.at<cv::Vec2f>(0, 3)[1]),
                               (float) base.rows);

        cv::Matx33f T = cv::Matx33f::zeros();
        T(0, 0) = 1;
        T(1, 1) = 1;
        T(2, 2) = 1;
        T(0, 2) = -x_start;
        T(1, 2) = -y_start;

        T = T * H;
        cv::Mat panorama;
        cv::warpPerspective(attach, panorama, T, cv::Size(x_end - x_start + 1, y_end - y_start + 1), cv::INTER_LINEAR);

        cv::Mat roi(panorama, cv::Rect(-x_start, -y_start, base.cols, base.rows));
        base.copyTo(roi, base);

        return panorama;
    }


    cv::Mat
    drawMatches(const cv::Mat &img1, const cv::Mat &img2, std::vector<cv::Vec3f> &p1, std::vector<cv::Vec3f> &p2) {

        cv::Mat_<cv::Vec3b> img1Converted;
        img1.convertTo(img1Converted, CV_8UC3);
        cv::Mat_<cv::Vec3b> img2Converted;
        img2.convertTo(img2Converted, CV_8UC3);

        cv::Mat_<cv::Vec3b> image(img1.rows + img2.rows, std::max(img1.cols, img2.cols));
        img1Converted.copyTo(image(cv::Rect(0, 0, img1.cols, img1.rows)));
        img2Converted.copyTo(image(cv::Rect(0, img1.rows, img2.cols, img2.rows)));


        if (p1.size() != p2.size())
            throw std::runtime_error("Mismatched sizes for matched point arrays!");

        for (unsigned i = 0; i < p1.size(); i++) {
            cv::line(image, {static_cast<int>(p1[i][0] / p1[i][2]), static_cast<int>(p1[i][1] / p1[i][2])},
                     {static_cast<int>(p2[i][0] / p2[i][2]), static_cast<int>(img1.rows + p2[i][1] / p2[i][2])},
                     cv::Scalar(0, 0, 255), 1);
        }

        for (unsigned i = 0; i < p1.size(); i++) {
            cv::circle(image, {static_cast<int>(p1[i][0] / p1[i][2]), static_cast<int>(p1[i][1] / p1[i][2])}, 2,
                       cv::Scalar(0, 255, 0), 2);
            cv::circle(image,
                       {static_cast<int>(p2[i][0] / p2[i][2]), static_cast<int>(img1.rows + p2[i][1] / p2[i][2])}, 2,
                       cv::Scalar(0, 255, 0), 2);
        }
        return image;
    }


    RawOrbMatches extractRawOrbMatches(const cv::Mat &img1, const cv::Mat &img2) {
        cv::Ptr<cv::ORB> orb = cv::ORB::create();
        orb->setMaxFeatures(30000);
        cv::Ptr<cv::BFMatcher> matcher = cv::BFMatcher::create(cv::NORM_HAMMING);

        std::vector<cv::KeyPoint> keypoints1, keypoints2;
        cv::Mat descriptors1, descriptors2;

        orb->detectAndCompute(img1, cv::noArray(), keypoints1, descriptors1);
        orb->detectAndCompute(img2, cv::noArray(), keypoints2, descriptors2);

        RawOrbMatches result;
        result.keypoints1.reserve(keypoints1.size());
        for (const auto &kp: keypoints1)
            result.keypoints1.push_back({kp.pt.x, kp.pt.y, 1.0f});

        result.keypoints2.reserve(keypoints2.size());
        for (const auto &kp: keypoints2)
            result.keypoints2.push_back({kp.pt.x, kp.pt.y, 1.0f});


        std::vector<std::vector<cv::DMatch>> matches;

        matcher->knnMatch(descriptors1, descriptors2, matches, 2);
        for (const auto &m: matches) {
            RawOrbMatches::Match match;
            match.closest = m[0].trainIdx;
            match.closestDistance = m[0].distance;
            match.secondClosest = m[1].trainIdx;
            match.secondClosestDistance = m[1].distance;
            result.matches_1_2[m[0].queryIdx] = match;
        }

        matcher->knnMatch(descriptors2, descriptors1, matches, 2);
        for (const auto &m: matches) {
            RawOrbMatches::Match match;
            match.closest = m[0].trainIdx;
            match.closestDistance = m[0].distance;
            match.secondClosest = m[1].trainIdx;
            match.secondClosestDistance = m[1].distance;
            result.matches_2_1[m[0].queryIdx] = match;
        }

        return result;
    }


}
